import libvirt
import xml.etree.ElementTree as ET
import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import base64
import subprocess

class Virtual:
    def __init__(self):
        self.conn = libvirt.open('qemu:///system')
        # AES 加密密钥，建议在实际使用中从安全配置中获取
        self.aes_key = b'e7278808408f4095bc4003a972f07ad6'  # 32 bytes for AES-256
        self.aes_iv = b'9de2c17f4d1aefff'  # 16 bytes for AES

    def get_kvm_networks(self):
        networks = []
        try:
            for net in self.conn.listAllNetworks():
                network_info = {
                    'name': net.name(),
                    'uuid': net.UUIDString(),
                    'active': net.isActive()
                }
                networks.append(network_info)
            return networks
        except Exception as e:
            print(f"获取 KVM 网络接口信息时出错: {e}")
            return []

    def _encrypt_port(self, port):
        """使用 AES 加密端口号并添加随机字符前缀，最后进行 Base64 编码"""
        if port is None:
            return None
        try:
            # 生成8位随机字符
            import random
            import string
            random_prefix = ''.join(random.choices(string.ascii_letters + string.digits, k=8)).lower()
            
            cipher = AES.new(self.aes_key, AES.MODE_CBC, self.aes_iv)
            encrypted = cipher.encrypt(pad(str(port).encode(), AES.block_size))
            encrypted_str = base64.b64encode(encrypted).decode()
            # 对结果进行 Base64 编码
            return base64.b64encode(f"{random_prefix}:{encrypted_str}".encode()).decode()
        except Exception as e:
            print(f"端口加密失败: {e}")
            return None

    def _decrypt_port(self, encrypted_port):
        """先进行 Base64 解码，然后使用 AES 解密带随机字符前缀的端口号"""
        if encrypted_port is None:
            return None
        try:
            # 先进行 Base64 解码
            decoded_port = base64.b64decode(encrypted_port).decode()
            
            # 验证输入格式
            if ':' not in decoded_port:
                raise ValueError("Invalid encrypted port format")
                
            # 分割随机前缀和加密内容
            _, encrypted_content = decoded_port.split(':', 1)
            
            cipher = AES.new(self.aes_key, AES.MODE_CBC, self.aes_iv)
            decrypted = cipher.decrypt(base64.b64decode(encrypted_content))
            return int(unpad(decrypted, AES.block_size).decode())
        except Exception as e:
            print(f"端口解密失败: {e}")
            return None


    def get_all_vms(self):
        vms = []
        for vm in self.conn.listAllDomains():
            # 获取虚拟机的 XML 描述
            xml_desc = vm.XMLDesc()
            # 获取虚拟机的状态
            state, _ = vm.state()

            # 解析 XML 以获取 graphics 类型
            root = ET.fromstring(xml_desc)
            graphics_list = []
            for graphics in root.findall('.//graphics'):
                graphics_type = graphics.get('type')
                # 只处理 spice 和 vnc 类型
                if graphics_type not in ['spice', 'vnc']:
                    continue
                    
                port = graphics.get('port')
                if port and port != '-1':
                    graphics_port = self._encrypt_port(int(port))
                else:
                    graphics_port = None
                
                graphics_list.append({
                    'type': graphics_type,
                    'ep': graphics_port
                })

            # 获取内存信息
            memory = root.find('memory')
            memory_value = int(memory.text) if memory is not None else None

            # 获取 CPU 信息
            vcpu = root.find('vcpu')
            vcpu_count = int(vcpu.text) if vcpu is not None else None

            # 获取磁盘信息
            disks = []
            for disk in root.findall('.//disk[@device="disk"]'):
                source = disk.find('source')
                if source is not None:
                    file_path = source.get('file')
                    if file_path:
                        try:
                            disk_size = os.path.getsize(file_path)
                        except FileNotFoundError:
                            disk_size = None
                        disk_info = {
                            'path': file_path,
                            'size': disk_size
                        }
                        disks.append(disk_info)

            # 构造包含虚拟机信息的字典
            vm_info = {
                'name': vm.name(),
                'id': vm.ID(),
                'uuid': vm.UUIDString(),
                'state': state,
                'xml_desc': xml_desc,
                'graphics': graphics_list,  # 修改为 graphics 列表
                'memory': memory_value,
                'vcpu_count': vcpu_count,
                'disks': disks
            }
            vms.append(vm_info)
        return vms

    def get_vm_state(self, vm_name):
        return self.conn.lookupByName(vm_name).state()

    def start_vm(self, vm_name):
        """
        启动指定名称的虚拟机，自动处理暂停状态的恢复
        :param vm_name: 要启动的虚拟机的名称
        :return: 如果虚拟机成功启动，返回 True；如果启动失败，返回 False
        """
        try:
            vm = self.conn.lookupByName(vm_name)
            state, _ = vm.state()
            # 如果虚拟机处于暂停状态，则恢复
            if state == libvirt.VIR_DOMAIN_PAUSED:
                return self.resume_vm(vm_name)
            # 如果虚拟机已经在运行，返回 False
            if vm.isActive():
                print(f"虚拟机 {vm_name} 已经在运行。")
                return False
            # 否则正常启动
            vm.create()
            print(f"虚拟机 {vm_name} 已成功启动。")
            return True
        except libvirt.libvirtError as e:
            print(f"启动虚拟机 {vm_name} 时出错: {e}")
            return False

    def resume_vm(self, vm_name):
        """
        恢复已暂停的虚拟机
        :param vm_name: 虚拟机名称
        :return: 操作是否成功
        """
        vm = self.conn.lookupByName(vm_name)
        if vm is None:
            return False
        try:
            vm.resume()
            print(f"虚拟机 {vm_name} 已从暂停状态恢复。")
            return True
        except libvirt.libvirtError as e:
            print(f"恢复虚拟机时出错: {e}")
            return False

    def shutdown_vm(self, vm_name):
        """
        关闭指定名称的虚拟机。

        :param vm_name: 要关闭的虚拟机的名称
        :return: 如果虚拟机成功关闭，返回 True；如果虚拟机已经关闭或关闭失败，返回 False
        """
        try:
            vm = self.conn.lookupByName(vm_name)
            if not vm.isActive():
                print(f"虚拟机 {vm_name} 已经关闭。")
                return False
            vm.shutdown()
            print(f"虚拟机 {vm_name} 已成功关闭。")
            return True
        except libvirt.libvirtError as e:
            print(f"关闭虚拟机 {vm_name} 时出错: {e}")
            return False

    def force_shutdown_vm(self, vm_name):
        """
        强制关闭指定虚拟机
        """
        try:
            # 获取虚拟机对象
            vm = self.conn.lookupByName(vm_name)
            # 调用 destroy 方法强制关闭
            vm.destroy()
            return True
        except libvirt.libvirtError as e:
            print(f"强制关闭虚拟机 {vm_name} 时出错: {e}")
            return False

    def clone_vm(self, source_vm_name, clone_vm_name):
        """
        克隆指定名称的虚拟机，并使用 qemu-img 创建新磁盘，指定源磁盘为后端磁盘。

        :param source_vm_name: 源虚拟机的名称
        :param clone_vm_name: 克隆后虚拟机的名称
        :return: 如果克隆成功，返回 True；如果克隆失败，返回 False
        """
        try:
            source_vm = self.conn.lookupByName(source_vm_name)
            source_xml = source_vm.XMLDesc()
            root = ET.fromstring(source_xml)

            # 修改虚拟机名称
            root.find('name').text = clone_vm_name

            # 删除 UUID 元素
            uuid_element = root.find('uuid')
            if uuid_element is not None:
                root.remove(uuid_element)

            # 删除所有网卡的 MAC 地址
            for interface in root.findall('.//interface'):
                mac_element = interface.find('mac')
                if mac_element is not None:
                    interface.remove(mac_element)
                # 删除网络相关的唯一标识，如 bridge 名称
                bridge_element = interface.find('source[@bridge]')
                if bridge_element is not None:
                    bridge_element.set('bridge', '')

            # 删除设备序列号
            for disk in root.findall('.//disk[@device="disk"]'):
                serial_element = disk.find('serial')
                if serial_element is not None:
                    disk.remove(serial_element)

            # 处理多个磁盘
            for disk in root.findall('.//disk[@device="disk"]'):
                source = disk.find('source')
                if source is not None:
                    file_path = source.get('file')
                    if file_path:
                        dir_path = os.path.dirname(file_path)
                        file_name = os.path.basename(file_path)
                        file_name_without_ext, file_ext = os.path.splitext(file_name)
                        # 新磁盘名等于旧的+虚拟机名字
                        new_file_name = f"{file_name_without_ext}-{clone_vm_name}{file_ext}"
                        new_file_path = os.path.join(dir_path, new_file_name)
                        source.set('file', new_file_path)

                        # 使用 qemu-img 检测源磁盘格式
                        import subprocess
                        try:
                            result = subprocess.run(['qemu-img', 'info', '--output=json', file_path], capture_output=True, text=True, check=True)
                            import json
                            info = json.loads(result.stdout)
                            backing_format = info.get('format')
                            if not backing_format:
                                print(f"无法检测到源磁盘 {file_path} 的格式。")
                                return False
                        except subprocess.CalledProcessError as e:
                            print(f"检测源磁盘 {file_path} 格式时出错: {e}")
                            return False

                        # 使用 qemu-img 创建新磁盘，指定源磁盘为后端磁盘
                        try:
                            subprocess.run(['qemu-img', 'create', '-f', 'qcow2', '-F', backing_format, '-b', file_path, new_file_path], check=True)
                        except subprocess.CalledProcessError as e:
                            print(f"创建新磁盘 {new_file_path} 时出错: {e}")
                            return False

                        # 修改 XML 中的磁盘格式为 qcow2
                        driver = disk.find('driver')
                        if driver is not None:
                            driver.set('type', 'qcow2')
                        else:
                            # 如果 driver 元素不存在，则创建一个
                            driver = ET.SubElement(disk, 'driver')
                            driver.set('name', 'qemu')
                            driver.set('type', 'qcow2')

            # 不再移除 bus 属性
            # 移除可能导致冲突的其他设备相关配置
            # for disk in root.findall('.//disk[@device="disk"]'):
            #     bus_element = disk.find('target[@bus]')
            #     if bus_element is not None:
            #         bus_element.set('bus', '')

            new_xml = ET.tostring(root, encoding='unicode')

            # 克隆虚拟机
            clone = self.conn.defineXML(new_xml)
            if clone:
                print(f"虚拟机 {clone_vm_name} 克隆成功。")
                return True
            else:
                print(f"虚拟机 {clone_vm_name} 克隆失败。")
                return False
        except libvirt.libvirtError as e:
            print(f"克隆虚拟机 {source_vm_name} 时出错: {e}")
            return False

    def create_vm(self, xml_desc):
        """
        根据提供的 XML 描述创建虚拟机。

        :param xml_desc: 虚拟机的 XML 描述字符串
        :return: 如果虚拟机创建成功，返回 True；如果创建失败，返回 False
        """
        try:
            vm = self.conn.defineXML(xml_desc)
            if vm:
                print("虚拟机创建成功。")
                return True
            else:
                print("虚拟机创建失败。")
                return False
        except libvirt.libvirtError as e:
            print(f"创建虚拟机时出错: {e}")
            return False

    def get_vm_graphics_port(self, vm_name, encrypted=False):
        """
        获取指定名称虚拟机的 graphics 端口，支持 VNC 和 Spice 协议。

        :param vm_name: 要获取 graphics 端口的虚拟机的名称
        :param encrypted: 是否返回加密后的端口
        :return: 如果找到端口，返回端口号；如果未找到或出现错误，返回 None
        """
        try:
            vm = self.conn.lookupByName(vm_name)
            xml_desc = vm.XMLDesc()
            root = ET.fromstring(xml_desc)
            # 查找 VNC 或 Spice 类型的 graphics 元素
            graphics = root.find('.//graphics[@type="vnc"]') or root.find('.//graphics[@type="spice"]')
            if graphics is not None:
                port = graphics.get('port')
                if port and port != '-1':
                    port = int(port)
                    return self._encrypt_port(port) if encrypted else port
        except libvirt.libvirtError as e:
            print(f"获取虚拟机 {vm_name} 的 graphics 端口时出错: {e}")
        return None

    def decrypt_graphics_port(self, encrypted_port):
        """
        解密加密的 graphics 端口

        :param encrypted_port: 加密后的端口字符串
        :return: 如果解密成功，返回端口号；如果解密失败，返回 None
        """
        return self._decrypt_port(encrypted_port)

    def suspend_vm(self, vm_name):
        """
        暂停指定虚拟机
        :param vm_name: 虚拟机名称
        :return: 操作是否成功
        """
        vm = self.conn.lookupByName(vm_name)
        if vm is None:
            return False
        try:
            # 直接暂停虚拟机
            vm.suspend()
            return True
        except libvirt.libvirtError as e:
            print(f"暂停虚拟机时出错: {e}")
            return False

    def get_libvirt_images(self):
        """
        获取 /var/lib/libvirt 目录下镜像的详细信息
        """
        try:
            image_dir = '/var/lib/libvirt'
            images = []
            for root, dirs, files in os.walk(image_dir):
                for file in files:
                    file_path = os.path.join(root, file)
                    try:
                        # 使用 qemu-img info 命令获取镜像信息
                        result = subprocess.run(['qemu-img', 'info', '--output=json', file_path], capture_output=True, text=True)
                        if result.returncode == 0:
                            info = eval(result.stdout)
                            # 获取镜像格式
                            format = info.get('format')
                            # 获取定义的存储大小
                            virtual_size = info.get('virtual-size')
                            # 获取已使用大小
                            actual_size = info.get('actual-size')
                            # 检查是否指定后端磁盘
                            backing_file = info.get('backing-filename')
                            has_backing_file = backing_file is not None

                            file_info = {
                                'name': file,
                                'path': file_path,
                                'size': os.path.getsize(file_path),
                                'format': format,
                                'virtual_size': virtual_size,
                                'actual_size': actual_size,
                                'has_backing_file': has_backing_file
                            }
                            images.append(file_info)
                    except Exception as e:
                        print(f"获取 {file_path} 信息时出错: {e}")

            return images
        except Exception as e:
            print(f"获取镜像信息时出错: {e}")
            return []
    def delete_vm(self, vm_name, disks_to_delete):
        try:
            vm = self.conn.lookupByName(vm_name)
            if vm.isActive():
                vm.destroy()  # 如果虚拟机正在运行，先关闭它
            vm_xml = vm.XMLDesc()
            root = ET.fromstring(vm_xml)
            valid_disks = []
            for disk in root.findall('.//disk[@device="disk"]'):
                source = disk.find('source')
                if source is not None:
                    file_path = source.get('file')
                    if file_path and file_path.startswith('/var/lib/libvirt/images'):
                        valid_disks.append(file_path)

            for disk in disks_to_delete:
                if disk in valid_disks:
                    try:
                        os.remove(disk)
                        print(f"成功删除磁盘: {disk}")
                    except Exception as e:
                        print(f"删除磁盘 {disk} 时出错: {e}")
            vm.undefine()  # 从 libvirt 中删除虚拟机定义
            return True
        except Exception as e:
            print(f"删除虚拟机 {vm_name} 时出错: {e}")
            return False
